
#include "webrtc_aec3.h"

#include "modules/audio_processing/aec3/echo_canceller3.h"
#include "modules/audio_processing/ns/noise_suppressor.h"
#include "modules/audio_processing/gain_controller2.h"
#include "modules/audio_processing/audio_buffer.h"

using namespace webrtc;

typedef struct 
{
    uint32_t    options;
    uint32_t	samples_per_frame;
    uint32_t    clock_rate;
    uint32_t	channel_count;
    uint32_t	frame_length;
    uint32_t	num_bands;

    EchoControl     *aec;
    NoiseSuppressor *ns;
    GainController2 *agc;
    AudioBuffer     *cap_buf;
    AudioBuffer     *rend_buf;
} webrtc_aec3_entity_t;

int webrtc_aec3_create(void** pp_entity,
                      uint32_t clock_rate,
                      uint32_t channel_count,
                      uint32_t samples_per_frame,
                      uint32_t tail_ms,
                      uint32_t options)
{
    // check the argument
    if(!pp_entity) {
        return WEBRTC_AEC_ERROR_INVALID_ARGUMENT;
    }

    if (clock_rate != 16000 && clock_rate != 32000 && clock_rate != 48000) {
    	return WEBRTC_AEC_ERROR_UNSUPPORT_CLOCK_RATE;
    }

    webrtc_aec3_entity_t* p_entity = new webrtc_aec3_entity_t;

    p_entity->options = options;    
    p_entity->channel_count = channel_count;
    p_entity->samples_per_frame = samples_per_frame;
    p_entity->clock_rate = clock_rate;
    p_entity->frame_length = clock_rate/100;
    p_entity->num_bands = clock_rate/16000;
    
    p_entity->aec = new EchoCanceller3(EchoCanceller3Config(), clock_rate,
    				   channel_count, channel_count);
    
    p_entity->cap_buf = new AudioBuffer(clock_rate, channel_count, clock_rate,
                        	    channel_count, clock_rate, channel_count);
    p_entity->rend_buf = new AudioBuffer(clock_rate, channel_count, clock_rate,
                       		     channel_count, clock_rate, channel_count);

    if (options & WEBRTC_AEC_OPTION_USE_NOISE_SUPPRESSOR) {
        NsConfig cfg;
        // Valid values are 6, 12, 18, 21 dB
        cfg.target_level = NsConfig::SuppressionLevel::k12dB;
        p_entity->ns = new NoiseSuppressor(cfg, clock_rate, channel_count);
    }

    if (options & WEBRTC_AEC_OPTION_USE_GAIN_CONTROLLER) {
        p_entity->agc = new GainController2();
        p_entity->agc->Initialize(clock_rate);
	
        AudioProcessing::Config::GainController2 cfg;
        cfg.adaptive_digital.enabled = true;
        if (GainController2::Validate(cfg)) {
            p_entity->agc->ApplyConfig(cfg);
        }
    }

    *pp_entity = (void*)p_entity;
    return WEBRTC_AEC_ERROR_SUCCESS;
}

int webrtc_aec3_destroy(void* p_entity)
{
    webrtc_aec3_entity_t* entity_ptr = (webrtc_aec3_entity_t*)p_entity;

    if(!p_entity)
        return WEBRTC_AEC_ERROR_INVALID_ARGUMENT;
    
    if (entity_ptr->aec) {
    	delete entity_ptr->aec;
    	entity_ptr->aec = NULL;
    }
    if (entity_ptr->ns) {
    	delete entity_ptr->ns;
    	entity_ptr->ns = NULL;
    }
    if (entity_ptr->agc) {
    	delete entity_ptr->agc;
    	entity_ptr->agc = NULL;
    }
    
    if (entity_ptr->cap_buf) {
    	delete entity_ptr->cap_buf;
    	entity_ptr->cap_buf = NULL;
    }    
    if (entity_ptr->rend_buf) {
    	delete entity_ptr->rend_buf;
    	entity_ptr->rend_buf = NULL;
    }

    return WEBRTC_AEC_ERROR_SUCCESS;
}

int webrtc_aec3_cancel_echo(void* p_entity,
					    int16_t *recv_frame,
					    const int16_t *play_frame)
{
    webrtc_aec3_entity_t *entity_ptr = (webrtc_aec3_entity_t*)p_entity;
    int32_t i;

    if(!entity_ptr && !recv_frame && !play_frame)
        return WEBRTC_AEC_ERROR_INVALID_ARGUMENT; 

    for (i = 0; i < entity_ptr->samples_per_frame; i += entity_ptr->frame_length) {

        StreamConfig scfg(entity_ptr->clock_rate, entity_ptr->channel_count);

    	entity_ptr->cap_buf->CopyFrom(recv_frame + i, scfg);
    	entity_ptr->rend_buf->CopyFrom(play_frame + i, scfg);

    	if (entity_ptr->clock_rate > 16000) {
      	    entity_ptr->cap_buf->SplitIntoFrequencyBands();
      	    entity_ptr->rend_buf->SplitIntoFrequencyBands();
    	}

    	entity_ptr->aec->AnalyzeCapture(entity_ptr->cap_buf);
      	entity_ptr->aec->AnalyzeRender(entity_ptr->rend_buf);
      	
      	if (entity_ptr->ns) {
      	    entity_ptr->ns->Analyze(*entity_ptr->cap_buf);
      	    entity_ptr->ns->Process(entity_ptr->cap_buf);
      	}
      	
      	entity_ptr->aec->ProcessCapture(entity_ptr->cap_buf, false);

      	if (entity_ptr->agc) {
      	    entity_ptr->agc->Process(entity_ptr->cap_buf);
      	}

    	if (entity_ptr->clock_rate > 16000) {
      	    entity_ptr->cap_buf->MergeFrequencyBands();
	}

     	entity_ptr->cap_buf->CopyTo(scfg, recv_frame + i);
    }

    return WEBRTC_AEC_ERROR_SUCCESS;
}